import { Annotation } from '@langchain/langgraph';
import { baseInputVars, LLMNode, LLMProvider } from "../../../core";
import { Jimp } from "jimp";
import { IEditItem } from '../data';

export class IdentifyImgObjectNode extends LLMNode {
    name = "IdentifyImgObjectNode";
    inputVars = Annotation.Root({
        ...baseInputVars.spec,
        imageBuffer: Annotation<any>,
        image: Annotation<string>,
        color: Annotation<string>,

        editList: Annotation<IEditItem[]>,
    });

    outputVars = Annotation.Root({
    });

    private drawBBox(image, x1, y1, width, height, color: "blue" | "red" | "white") {
        // 根据颜色参数确定RGB值
        let red, green, blue;
        switch (color) {
            case "blue":
                red = 0;
                green = 0;
                blue = 255;
                break;
            case "red":
                red = 255;
                green = 0;
                blue = 0;
                break;
            case "white":
                red = 255;
                green = 255;
                blue = 255;
                break;
            default:
                red = 255;
                green = 0;
                blue = 0;
        }

        // 定义绘制线条的辅助函数
        const drawLine = (startX, startY, lineWidth, lineHeight) => {
            image.scan(startX, startY, lineWidth, lineHeight, (x, y, idx) => {
                image.bitmap.data[idx] = red;
                image.bitmap.data[idx + 1] = green;
                image.bitmap.data[idx + 2] = blue;
                image.bitmap.data[idx + 3] = 255;
            });
        };

        // 使用scan方法绘制矩形框的四条边
        drawLine(x1, y1, width, 1);                           // 上边框
        drawLine(x1, y1 + height - 1, width, 1);              // 下边框
        drawLine(x1, y1, 1, height);                          // 左边框
        drawLine(x1 + width - 1, y1, 1, height);

        return image;
    }

    private async genObjectInfo(editItem: IEditItem) {
        let image = await Jimp.read(this.inputVars.State.image);
        let [x1, y1, x2, y2] = editItem.bbox.map(x => parseFloat(x.toFixed(2)));
        // 计算矩形的宽度和高度
        const width = x2 - x1;
        const height = y2 - y1;

        image = await this.drawBBox(image, x1, y1, width, height, this.inputVars.State.color as any);
        // await image.write(`${editItem.objID}.jpg`);
        const imageBase64 = await image.getBase64("image/jpeg");

        let colorKey = "红色";
        switch (this.inputVars.State.color) {
            case "blue":
                colorKey = "蓝色";
                break;
            case "red":
                colorKey = "红色";
                break;
            case "white":
                colorKey = "绿色";
                break;
        }

        const model = LLMProvider.current.DoubaoSeed16Vision;
        model.thinking = "disabled";
        const prompts = [
            {
                role: "user",
                content: [
                    {
                        type: "image_url",
                        image_url: {
                            url: imageBase64
                        }
                    },
                    {
                        text: `识别定位${colorKey}bbox
        使用json的方式输出信息,
        {
            category: "主体信息（人物、身体部位、物品、物品部件、文字）"，
            posization: "主体相对于整图的位置描述、不要包含bbox的信息",
            description: "主体详细描述，如果是文字的话描述文字字体、大小、颜色等",
        }`,
                        type: "text"
                    }
                ]
            }
        ] as any;
        const outputStream = await model.stream(prompts);
        for await (const chunk of outputStream) {
        }

        this.usage.add(model.totalUsage);

        return {
            objID: editItem.objID,
            bboxObject: JSON.parse(model.lastFullResponse.replace(/\n/g, "").trim())
        }
    }

    async Run() {
        let plist = [];
        for (let item of this.inputVars.State.editList) {
            plist.push(this.genObjectInfo(item));
        }
        const result = await Promise.all(plist);

        this.outputVars.State[this.ResultField] = result;
        return this.outputVars.State;
    }
}